# -*- coding: utf-8 -*-
# 启动方式：mitmdump -s *.py --flow-detail 0 (不在窗口显示请求信息)
# 启动方式：mitmweb

import mitmproxy.http
from mitmproxy import ctx
import time
import datetime
import xlwt
import sqlite3

# 保存到数据库
conn = sqlite3.connect('http_connet.db')
# 只记录特定特征的网址
# url_key_word = "zljllcy.com"
# url_key_word = "*"
url_key_word = "hkitcloud.net"

class Saver:
    def __init__(self):
        self.num = 0
        self.requestNum = 0
        self.responseOrErrorNum = 0
        self.aa = 0
        # 特定请求域名才记录
        self.url_key_word = url_key_word
        # 是否记录
        self.mark_flag = 1
        self.all_arr = [[]]
        # 链接记录的字典
        self.connetc_dict = {}
        self.flow_tmp = []
        self.request_time = ""
        # 请求记录
        c = conn.cursor()
        c.execute('''CREATE TABLE IF NOT EXISTS httrequest
               (ID  INTEGER PRIMARY KEY  NOT NULL,
               URL           TEXT      NOT NULL,
               HOST          TEXT      NOT NULL,
               Path          TEXT      NOT NULL,
               RequestType   TEXT      ,
               RequestData   TEXT      ,
               ResponseData  TEXT      ,
               RecodTime     CHAR(25)  ,
               RequestTime   CHAR(25));''')# '请求路径', '请求域名', '请求path', '请求类型', '请求报文' ,'响应报文', '记录时间'
        c.close()
        conn.commit()

    def http_connect(self, flow: mitmproxy.http.HTTPFlow):
        """
        收到了来自客户端的 HTTP CONNECT 请求，此时只是客户端到mitmproxy的请求，还没触发到server的请求，所以此时还不会触发request等事件
        可以在此时对请求地址之类的做操作
        """
        self.flow_tmp = []
        tmp = [str(datetime.datetime.now()), '||', 'connect', '||', str(flow.id), '||', flow.type]
        self.log_out(tmp)
        self.connetc_dict['cid'] = str(flow.id)


    def request(self, flow: mitmproxy.http.HTTPFlow):
        """
        自客户端的 HTTP 请求被成功【完整】读取。
        """
        tmp = [str(datetime.datetime.now()), '||', 'request', '||', str(flow.id), '||', flow.request.url]
        self.log_out(tmp)
        self.mark_flag = 1
        # 请求的不是要抓取的地址
        if self.url_key_word not in flow.request.host:
            self.mark_flag = 0
        # 请求的是 js脚本、图片、音频等信息
        rq_type = flow.request.path.split('.')[-1]
        if rq_type in ['js', 'png', 'jpeg', 'jpg', 'gif', 'html', 'apk']:
            self.mark_flag = 0
        # 请求的类型是option的不记录
        if flow.request.method == 'OPTIONS':
            self.mark_flag = 0
        # 如果不过滤地址，则完整记录所有信息
        if url_key_word == "*":
            self.mark_flag = 1
        # unix时间戳转为本地时间

        if self.mark_flag == 0:
            # 删除字典中的记录
            self.connetc_dict.pop(str(flow.id))
        else:
            self.request_time = str(time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(flow.request.timestamp_start)))
            str_output = "-->" + self.request_time + "*****-->" + str(flow.request.path)
            ctx.log.info(str_output)
            # print("in cat req_url:" ,flow.request.host)
            self.num = self.num + 1
            self.requestNum = self.requestNum + 1
            # flow.start_time = time.time()
            self.flow_tmp = [flow.request.url, flow.request.host, flow.request.path]
            self.all_arr.append(self.flow_tmp)
            # 加入字典
            # 请求类型
            try:
                self.flow_tmp.append(flow.request.method)
            except:
                self.flow_tmp.append("")
            # 请求报文
            try:
                self.flow_tmp.append(str(flow.request.get_content(),'utf-8'))
            except:
                self.flow_tmp.append("记录请求报文出错")

    def error(self, flow: mitmproxy.http.HTTPFlow):
        """发生了一个 HTTP 错误。比如无效的服务端响应、连接断开等"""
        self.aa = self.aa + 1
        self.responseOrErrorNum = self.responseOrErrorNum + 1
        self.flow_tmp.append("Error response")

    def response(self, flow: mitmproxy.http.HTTPFlow):
        """自服务端端的 HTTP 响应被成功完整读取"""
        tmp = [str(datetime.datetime.now()), '||', 'response', '||', str(flow.id), '||', flow.request.url]
        self.log_out(tmp)
        if self.mark_flag == 1:
            self.aa = self.aa + 1
            self.responseOrErrorNum = self.responseOrErrorNum + 1
            flow.end_time = time.time()

            # 请求大小
            try:
                # self.flow_tmp.append(flow.request.headers['Content-Length'])
                pass
            except:
                # self.flow_tmp.append("")
                pass
            # 响应大小
            try:
                # self.flow_tmp.append(flow.response.headers['Content-Length'])
                pass
            except:
                # self.flow_tmp.append("")
                pass
            # 响应类型
            try:
                # self.flow_tmp.append(flow.response.headers['Content-Type'])
                pass
            except:
                # self.flow_tmp.append("")
                pass

            # 响应报文
            try:
                self.flow_tmp.append(flow.response.text)
            except:
                self.flow_tmp.append("")
            # 记录时间
            try:
                self.flow_tmp.append(time.strftime("%Y-%m-%d %H:%M:%S",time.localtime(time.time())))
            except:
                self.flow_tmp.append("")
            # 请求时间
            try:
                self.flow_tmp.append(self.request_time)
            except:
                self.flow_tmp.append("")
            # 请求响应时间差
            try:
                # time_gap = flow.end_time - flow.start_time
                # self.flow_tmp.append(time_gap)
                pass
            except Exception:
                # self.flow_tmp.append("")
                pass
            # self.formatoutput(flow)
            self.save(self.all_arr)

    def log_out(self, array):
        with open("mitmproxy.log", 'a+') as fp:
            try:
                fp.writelines(array)
                fp.write('\r\n')
            except Exception as bug:
                print("write file fail", bug)

    def save(self, array):
        # 是否保存到exce
        return 0
        save_2_excel = 0
        cor = conn.cursor()
        # print(array)
        # 保存到数据库
        try:
            cor.execute("INSERT INTO httrequest (ID,URL,HOST,Path,RequestType,RequestData,ResponseData,RecodTime,RequestTime) \
                        VALUES (?,?,?,?,?,?,?,?,?) ", \
            ((None, array[1][0], array[1][1], array[1][2], array[1][3], array[1][4], array[1][5], array[1][6],array[1][7])))
        except Exception as bug:
            ctx.log.info("save datadb faile")
            ctx.log.info(bug)
        conn.commit()
        if save_2_excel == 1:
            # 保存到文件
            workbook = xlwt.Workbook()
            worksheet = workbook.add_sheet('http_Request_Record')
            excel_heard = ['请求路径', '请求域名', '请求path', '请求类型', '请求报文' ,'响应报文', '记录时间']
            for i in range(len(excel_heard)):
                worksheet.write(0,i,excel_heard[i])
            for x in range(len(array)):
                for y in range(len(array[x])):
                    worksheet.write(x, y, array[x][y])
            workbook.save("http_requests.xls")


    def done(self):
        """
        addon 关闭或被移除，又或者 mitmproxy 本身关闭。由于会先等事件循环终止后再触发该事件，所以这是一个 addon 可以看见的最后一个事件。
        """
        conn.close()

